import pandas as pd


def f(data, index):
    """计算每支股票每天的连续5日、10日指标"""
    print(data.shape[0])
    print('开始计算。。。')
    # 连续5日成交量均值
    data['v_five'] = 0
    # 连续5日收盘价均值
    data['c_five'] = 0
    # 连续10日成交量均值
    data['v_ten'] = 0
    # 连续10日收盘价均值
    data['c_ten'] = 0

    for i in range(1, 11):
        """添加对应的成交量和收盘价的列，使得连续10行的数据在同一行"""
        data['a_v_{}'.format(i)] = data['量/万'].shift(i)
        data['a_c_{}'.format(i)] = data['收'].shift(i)

    for code in index:
        print(code)
        for i in range(0, data.shape[0]):
            if data['股票代码'][i] == code and data['天数'][i] >= 14:
                """因为前10个交易日无法计算连续10天的指标，所以从第11个交易日开始计算"""
                for n in range(1, 6):
                    data.loc[i, 'v_five'] = data['a_v_{}'.format(n)][i] + data['v_five'][i]
                    data.loc[i, 'c_five'] = data['a_c_{}'.format(n)][i] + data['c_five'][i]
                for n in range(1, 11):
                    data.loc[i, 'v_ten'] += data['a_v_{}'.format(n)][i]
                    data.loc[i, 'c_ten'] += data['a_c_{}'.format(n)][i]
    data['v_five'] = round(data['v_five']/5, 2)
    data['c_five'] = round(data['c_five'] / 5, 2)
    data['v_ten'] = round(data['v_ten'] / 10, 2)
    data['c_ten'] = round(data['c_ten'] / 10, 2)

    for i in range(1, 11):
        del data['a_v_{}'.format(i)]
        del data['a_c_{}'.format(i)]

    for i in range(0, data.shape[0]):
        """为了方便后期处理，用程序实现求每支股票每天交易量和收盘价与指标的差值"""
        if data['v_five'][i] != 0:
            data.loc[i, 'v_five_r'] = data['量/万'][i]-data['v_five'][i]
        if data['c_five'][i] != 0:
            data.loc[i, 'c_five_r'] = data['收'][i]-data['c_five'][i]
        if data['v_ten'][i] != 0:
            data.loc[i, 'v_ten_r'] = data['量/万'][i] - data['v_ten'][i]
        if data['c_ten'][i] != 0:
            data.loc[i, 'c_ten_r'] = data['收'][i] - data['c_ten'][i]
    print('开始写入。。。')
    data.to_csv('./four结果.csv', header=True, index=False, encoding="utf_8_sig")
    print('写入成功...')


if __name__ == '__main__':
    data = pd.read_csv('./four.csv')
    index = [600048, 600183, 600323, 600325, 600332, 600383, 600872
        , 601318, 28, 49, 69, 513, 636, 651,
             921, 1914, 2027, 2035, 2060, 2138, 2152,
             2233, 2303, 2352, 2449, 2461, 2475, 2511,
             300014, 300115]
    f(data, index)

